package net.cyue.ort.llm.generator;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategies;
import com.fasterxml.jackson.databind.annotation.JsonNaming;

import net.cyue.ort.llm.config.ModelConfig;
import net.cyue.util.StreamUtil;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;

@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
public class GenerationConfig {
    private int maxNewTokens;
    @JsonProperty("bos_token_id")
    private Long bosTokenID;
    @JsonProperty("pad_token_id")
    private Long padTokenID;
    @JsonProperty("eos_token_id")
    private long[] eosTokenIDs;
    private boolean doSample;
    private float repetitionPenalty;
    private float temperature;
    private float topP;
    private int topK;
    private int numBeams;  // beam search的束宽
    private int beamDepth;
    private float lengthPenalty; // 长度惩罚因子
    private int ngramSize; // ngram的大小
    private float ngramPenalty; // ngram惩罚系数

    public GenerationConfig(
        int maxNewTokens,
        Long bosTokenID,
        Long padTokenID,
        long[] eosTokenIDs,
        boolean doSample,
        float repetitionPenalty,
        float temperature,
        float topP,
        int topK,
        int numBeams,
        int beamDepth,
        float lengthPenalty,
        int ngramSize,
        float ngramPenalty
    ) {
        this.maxNewTokens = maxNewTokens;
        this.bosTokenID = bosTokenID;
        this.padTokenID = padTokenID;
        this.eosTokenIDs = eosTokenIDs;
        this.doSample = doSample;
        this.repetitionPenalty = repetitionPenalty;
        this.temperature = temperature;
        this.topP = topP;
        this.topK = topK;
        this.numBeams = numBeams;
        this.beamDepth = beamDepth;
        this.lengthPenalty = lengthPenalty;
        this.ngramSize = ngramSize;
        this.ngramPenalty = ngramPenalty;
    }

    @Override
    public String toString() {
        return "GenerationConfig {" +
            "\n\tmaxNewTokens=" + maxNewTokens +
            ",\n\tbosTokenID=" + bosTokenID +
            ",\n\tpadTokenID=" + padTokenID +
            ",\n\teosTokenIDs=" + Arrays.toString(eosTokenIDs) +
            ",\n\tdoSample=" + doSample +
            ",\n\trepetitionPenalty=" + repetitionPenalty +
            ",\n\ttemperature=" + temperature +
            ",\n\ttopP=" + topP +
            ",\n\ttopK=" + topK +
            ",\n\tnumBeams=" + numBeams +
            ",\n\tlengthPenalty=" + lengthPenalty +
            ",\n\tngramSize=" + ngramSize +
            ",\n\tngramPenalty=" + ngramPenalty +
            "\n}";
    }

    public static GenerationConfig fromFile(String configPath) throws IOException {
        return GenerationConfig.fromInputStream(new FileInputStream(configPath));
    }
    public static GenerationConfig fromInputStream(InputStream stream) throws IOException {
        return GenerationConfig.fromContent(StreamUtil.readAsString(stream));
    }
    public static GenerationConfig fromContent(String content) throws IOException {
        ObjectMapper objectMapper = new ObjectMapper();
        JsonNode node = objectMapper.readTree(content);
        GenerationConfig config = GenerationConfig.defaultConfig();
        if (node.has("bos_token_id")) {
            config.bosTokenID = node.get("bos_token_id").asLong();
        }
        if (node.has("pad_token_id")) {
            config.padTokenID = node.get("pad_token_id").asLong();
        }
        JsonNode eosTokenIDsNode = node.get("eos_token_id");
        if (eosTokenIDsNode != null && eosTokenIDsNode.isArray()) {
            config.eosTokenIDs = objectMapper.readValue(eosTokenIDsNode.traverse(), new TypeReference<>() {});
        } else if (eosTokenIDsNode != null) {
            config.eosTokenIDs = new long[]{eosTokenIDsNode.asLong()};
        }
        if (node.has("do_sample")) {
            config.doSample = node.get("do_sample").asBoolean();
        }
        if (node.has("repetition_penalty")) {
            config.repetitionPenalty = Double.valueOf(node.get("repetition_penalty").asDouble()).floatValue();
        }
        if (node.has("temperature")) {
            config.temperature = Double.valueOf(node.get("temperature").asDouble()).floatValue();
        }
        if (node.has("top_p")) {
            config.topP = Double.valueOf(node.get("top_p").asDouble()).floatValue();
        }
        if (node.has("top_k")) {
            config.topK = node.get("top_k").asInt();
        }

        System.out.println(config);
        return config;
    }

    public static GenerationConfig defaultConfig() {
        GenerationConfig config = new GenerationConfig(
            1024,
            null,
            null,
            null,
            true,
            1.1f,
            0.8f,
            0.8f,
            20,
            3,  // 默认使用 beam search
            3,  // 默认 beam depth
            1.0f,
            3,     // 默认ngram大小为3
            1.5f   // 默认惩罚系数为1.5
        );
        System.out.println(config);
        return config;
    }

    public static GenerationConfig defaultConfig(ModelConfig modelConfig) {
        GenerationConfig config = new GenerationConfig(
            1024,
            modelConfig.bosTokenID,
            null,
            new long[]{modelConfig.eosTokenID},
            true,
            1.1f,
            0.8f,
            0.8f,
            20,
            3,  // 默认使用 beam search
            3,  // 默认 beam depth
            1.0f,
            3,     // 默认ngram大小为3
            1.5f   // 默认惩罚系数为1.5
        );
        System.out.println(config);
        return config;
    }

    public GenerationConfig setMaxNewTokens(int maxNewTokens) {
        this.maxNewTokens = maxNewTokens;
        return this;
    }
    public GenerationConfig setRepetitionPenalty(float repetitionPenalty) {
        this.repetitionPenalty = repetitionPenalty;
        return this;
    }
    public GenerationConfig setEosTokenID(long eosTokenID) {
        this.eosTokenIDs = new long[]{eosTokenID};
        return this;
    }
    public GenerationConfig setEosTokenIDs(long[] eosTokenIDs) {
        this.eosTokenIDs = eosTokenIDs;
        return this;
    }
    public GenerationConfig setNumBeams(int numBeams) {
        this.numBeams = numBeams;
        return this;
    }
    public GenerationConfig setBeamWidth(int beamWidth) {
        this.numBeams = beamWidth;
        return this;
    }
    public GenerationConfig setBeamDepth(int beamDepth) {
        this.beamDepth = beamDepth;
        return this;
    }
    public GenerationConfig setLengthPenalty(float lengthPenalty) {
        this.lengthPenalty = lengthPenalty;
        return this;
    }
    public GenerationConfig setNgramSize(int ngramSize) {
        this.ngramSize = ngramSize;
        return this;
    }
    public GenerationConfig setNgramPenalty(float ngramPenalty) {
        this.ngramPenalty = ngramPenalty;
        return this;
    }

    public int getMaxNewTokens() {
        return this.maxNewTokens;
    }
    public Long getBosTokenID() {
        return this.bosTokenID;
    }
    public Long getPadTokenID() {
        return this.padTokenID;
    }
    public long[] getEosTokenIDs() {
        return this.eosTokenIDs;
    }
    public boolean isDoSample() {
        return this.doSample;
    }
    public float getRepetitionPenalty() {
        return this.repetitionPenalty;
    }
    public float getTemperature() {
        return this.temperature;
    }
    public float getTopP() {
        return this.topP;
    }
    public int getTopK() {
        return this.topK;
    }
    public int getNumBeams() {
        return this.numBeams;
    }
    public int getBeamDepth() {
        return this.beamDepth;
    }
    public float getLengthPenalty() {
        return this.lengthPenalty;
    }
    public int getNgramSize() {
        return this.ngramSize;
    }
    public float getNgramPenalty() {
        return this.ngramPenalty;
    }


    // 生成配置的构建器
    public static class GenerationBuilder {
        private GenerationConfig config;
        public GenerationBuilder() {
            // 使用默认配置
            this.config = GenerationConfig.defaultConfig();
        }

        public GenerationBuilder setBoSTokenID(Long bosTokenID) {
            this.config.bosTokenID = bosTokenID;
            return this;
        }
        public GenerationBuilder setPadTokenID(Long padTokenID) {
            this.config.padTokenID = padTokenID;
            return this;
        }
        public GenerationBuilder setEosTokenID(Long eosTokenID) {
            this.config.eosTokenIDs = new long[]{eosTokenID};
            return this;
        }
        public GenerationBuilder setEosTokenIDs(long[] eosTokenIDs) {
            this.config.eosTokenIDs = eosTokenIDs;
            return this;
        }
        public GenerationBuilder setDoSample(boolean doSample) {
            this.config.doSample = doSample;
            return this;
        }
        public GenerationBuilder setRepetitionPenalty(float repetitionPenalty) {
            this.config.repetitionPenalty = repetitionPenalty;
            return this;
        }
        public GenerationBuilder setTemperature(float temperature) {
            this.config.temperature = temperature;
            return this;
        }
        public GenerationBuilder setTopP(float topP) {
            this.config.topP = topP;
            return this;
        }
        public GenerationBuilder setTopK(int topK) {
            this.config.topK = topK;
            return this;
        }
        public GenerationBuilder setNumBeams(int numBeams) {
            this.config.numBeams = numBeams;
            return this;
        }
        public GenerationBuilder setBeamWidth(int beamWidth) {
            this.config.numBeams = beamWidth;
            return this;
        }
        public GenerationBuilder setBeamDepth(int beamDepth) {
            this.config.beamDepth = beamDepth;
            return this;
        }
        public GenerationBuilder setLengthPenalty(float lengthPenalty) {
            this.config.lengthPenalty = lengthPenalty;
            return this;
        }
        public GenerationBuilder setNgramSize(int ngramSize) {
            this.config.ngramSize = ngramSize;
            return this;
        }
        public GenerationBuilder setNgramPenalty(float ngramPenalty) {
            this.config.ngramPenalty = ngramPenalty;
            return this;
        }

        public GenerationBuilder setConfig(GenerationConfig config) {
            this.config = config;
            return this;
        }

        public GenerationConfig build() {
            return this.config;
        }
    }
}
